import JT
import time

# Define which data to use for registration:
vert = 0
proj = 0
body = 1

Data = JT.Data.main()    
#===============================================================================
#    Setup DRR
#===============================================================================      
drr = JT.DRR()
drr.SetExternalCalibration(Data.Frames[vert].FixedImages[proj].calibration)

drr.AddVolume(Data.Frames[vert].Volumes[0].data)
drr._renWin.SetOffScreenRendering(0)

Testing = JT.Testing()

#===============================================================================
#    Setup Registration
#===============================================================================      
# Set Optimizer
gd = JT.RegularStepGradientDescentOptimizer()

reg=JT.MultiIterationRegistration()

reg.SetFixedImage(Data.Frames[vert].FixedImages[proj].data.GetImage(0))
reg.SetMovingImageGenerator(drr)
reg.SetRegionOfInterest(Data.Frames[vert].FixedImages[proj].ROIs[body].data)
reg.SetOptimizer(gd)
# set Metric
reg.SetImageMetric('bpc')

# Registration Options
reg._plot = True
reg._debug = True
reg._logGradientResult = False
gd._debug = False


#===============================================================================
# Setup Optimizer Parameters
#===============================================================================
# Set Multiresolution Parameters
reg.SetIterations(4)
reg.SetReductionFactor(0.5)

# Set the cost function
gd.SetCostFunction(reg)
# Set Scales 1/1000 for translations, 1 for rotations
scales = [1.0,1.0,1.0,1700.0,1700.0,1700.0]
reg.SetScales(scales)
# Modify Optimizer Parameters

gd.SetMinimumStepLength(.25)
gd._MaximumGradientStepLength = 20.0
gd._MinimumGradientStepLength = 0.25
gd.SetRelaxationFactor(0.75)
gd.SetNumberOfIterations(50)

def run(n_pose):    
    # Set the starting pose
    pose = Testing.GetStartPose(n_pose)
    print "Starting Pose: ", pose
    reg.SetInitialPose(pose)
    # For multi iteration registration, the max step size is reduced for 
    # each iteration according to the reductionFactor (0.5 by default)
    gd.SetMaximumStepLength(8)
    # Start the registration process
    reg.Start()

        

if __name__ == "__main__":
    run(n_pose=100)
    #time.sleep(7)
    JT.PlotResultsFromFile.main(0)


